import tensorflow as tf
import numpy as np

# f(x) = a*x**2 + b*x + c的导数

x = tf.Variable(0.0, name="x", dtype=tf.float32)
a = tf.constant(1.0)
b = tf.constant(-2.0)
c = tf.constant(1.0)

with tf.GradientTape() as tape:
    y = a * tf.pow(x, 2) + b * x + c

dy_dx = tape.gradient(y, x)
print(dy_dx)


# 对常量张量也可以求导，需要增加watch
with tf.GradientTape() as tape:
    tape.watch([a, b, c])
    y = a * tf.pow(x, 2) + b * x + c

dy_dx, dy_da, dy_db, dy_dc = tape.gradient(y, [x, a, b, c])
print(dy_da)
print(dy_dc)

# 可以求二阶导数
with tf.GradientTape() as tape2:
    with tf.GradientTape() as tape1:
        y = a*tf.pow(x,2) + b*x + c
    dy_dx = tape1.gradient(y,x)
dy2_dx2 = tape2.gradient(dy_dx,x)

print(dy2_dx2)


# 可以在autograph中使用

@tf.function
def f(x):
    a = tf.constant(1.0)
    b = tf.constant(-2.0)
    c = tf.constant(1.0)

    # 自变量转换成tf.float32
    x = tf.cast(x, tf.float32)
    with tf.GradientTape() as tape:
        tape.watch(x)
        y = a * tf.pow(x, 2) + b * x + c
    dy_dx = tape.gradient(y, x)

    return ((dy_dx, y))


tf.print(f(tf.constant(0.0)))
tf.print(f(tf.constant(1.0)))